package cn.itcast.tags.models.ml

import org.apache.spark.ml.clustering.{KMeans, KMeansModel}
import org.apache.spark.sql.{DataFrame, SparkSession}

object IrisClusterTest {
  def main(args: Array[String]): Unit = {
    val spark = SparkSession.builder()
      .appName(this.getClass.getSimpleName.stripSuffix("$"))
      .master("local[2]")
      .getOrCreate()
    import org.apache.spark.sql.functions._
    import spark.implicits._

    val irisDF: DataFrame = spark.read
      .format("libsvm")
      .load("datas/iris/iris_kmeans.txt")

    irisDF.printSchema()
    irisDF.show(10,false)

    val kMeans: KMeans = new KMeans()
      .setFeaturesCol("features")
      .setPredictionCol("prediction")
      .setK(3)
      .setMaxIter(20)

    val kMeansModel: KMeansModel = kMeans.fit(irisDF)

    kMeansModel.clusterCenters.foreach(println)

    val wssse: Double = kMeansModel.computeCost(irisDF)
    println(wssse)
    val preditionDF: DataFrame = kMeansModel.transform(irisDF)
    preditionDF.show(150,false)

    preditionDF.groupBy($"label",$"prediction")
    .count().show(20,false)
    spark.stop()
  }

}
